package bin_tree.heap.leetcode;

import java.util.Comparator;
import java.util.PriorityQueue;
import java.util.Queue;

/**
 * 面试题17.14，找到前k小元素
 * https://leetcode-cn.com/problems/smallest-k-lcci/
 * @author yuisama
 * @date 2022/01/07 21:26
 **/
public class TopK {
    public int[] smallestK(int[] arr, int k) {
        int[] ret = new int[k];
        // 边界条件注意
        if (arr.length == 0 || k == 0) {
            return ret;
        }
        // 构建一个最大堆
        // JDK的优先级队列是最小堆，如何把最小堆变为最大堆？
        // 就要用到我们比较器
        // 元素越小优先级越高，变为最大堆，元素越大优先级越高
        Queue<Integer> queue = new PriorityQueue<>(new Comparator<Integer>() {
            @Override
            public int compare(Integer o1, Integer o2) {
                return o2 - o1;
            }
        });
        // 遍历原数组，扫描元素，将元素加入最大堆
        for(int value : arr) {
            if (queue.size() < k) {
                queue.offer(value);
            }else {
                if (value < queue.peek()) {
                    queue.poll();
                    queue.offer(value);
                }
            }
        }
        // 当扫描完整个原数组时，queue中就存储了前k小个元素，依次出队即可
        for (int i = 0; i < k; i++) {
            ret[i] = queue.poll();
        }
        return ret;
    }
}